#include <iostream>
#include <random>
#include "backend/vertex_inverse_depth.h"
#include "backend/vertex_pose.h"
#include "backend/edge_reprojection.h"
#include "backend/edge_prior.h"
#include "backend/problem.h"

using namespace myslam::backend;
using namespace std;

#define RESET    1
#define ADDPRIOR 1

int window_size = 3;
int featureNums = 20;  // 每次产生 20 个新的特征被所有的帧观测到
/*
 * Frame : 保存每帧的姿态和观测
 */
struct Frame {
    Frame(Eigen::Matrix3d R, Eigen::Vector3d t) : Rwc(R), qwc(R), twc(t) {};
    Eigen::Matrix3d Rwc;
    Eigen::Quaterniond qwc;
    Eigen::Vector3d twc;

    unordered_map<int, Eigen::Vector3d> featurePerId; // 该帧观测到的特征以及特征id
};

struct Point3d {
    Point3d(Eigen::Vector3d p) : xyz(p) {
        start_frame_id = -1;
        featureId = -1;
        obs_cnt = 0;
    };

    Eigen::Vector3d xyz;
    Eigen::Vector3d xyz_in_startframe;
    int start_frame_id;
    int obs_cnt;
    int featureId;
    shared_ptr<VertexInverseDepth> verterx;
};

/*
 * 产生世界坐标系下的虚拟数据: 特征点, 以及每帧观测
 */
void GetSimDataInWordFrame(vector<Frame> &cameraPoses, vector<Point3d> &points) {

    int poseNums = 10;     // 相机数目
    double radius = 8;
    for (int n = 0; n < poseNums; ++n) {
        double theta = n * 2 * M_PI / (poseNums * 4); // 1/4 圆弧
        // 绕 z轴 旋转
        Eigen::Matrix3d R;
        R = Eigen::AngleAxisd(theta, Eigen::Vector3d::UnitZ());
        Eigen::Vector3d t = Eigen::Vector3d(radius * cos(theta) - radius, radius * sin(theta), 1 * sin(2 * theta));
        cameraPoses.push_back(Frame(R, t));
    }

    // 随机数生成三维特征点
    std::default_random_engine generator;
    std::normal_distribution<double> noise_pdf(0., 2. / 500.);  // 2pixel / focal
    int featureId = 0;
    // 控制滑动窗口次数
    for (int swf = 0; swf < poseNums - window_size; swf++) {
        // 滑动窗口每滑动一下，最新窗口帧上产生新的观测，其他最老帧看不到这个特征
        for (int j = 0; j < featureNums; ++j) {

            std::uniform_real_distribution<double> xy_rand(-4, 4.0);
            std::uniform_real_distribution<double> z_rand(4., 8.);

            Eigen::Vector3d Pw(xy_rand(generator), xy_rand(generator), z_rand(generator));
            Point3d pt(Pw);
            pt.start_frame_id = swf;
            pt.featureId = featureId;

            // 控制只有窗口内的相机看到这个特征
            for (size_t i = swf; i < swf + window_size; ++i) {
                Eigen::Vector3d Pc = cameraPoses[i].Rwc.transpose() * (Pw - cameraPoses[i].twc);
                pt.obs_cnt++;
                if (i == swf)
                    pt.xyz_in_startframe = Pc;

                Pc = Pc / Pc.z();  // 归一化图像平面
                Pc[0] += noise_pdf(generator);
                Pc[1] += noise_pdf(generator);
                cameraPoses[i].featurePerId.insert(make_pair(featureId, Pc));
            }

            points.push_back(pt);
            featureId++;
        }
    }
}

int main() {
    // 准备数据
    vector<Frame> cameras;
    vector<Point3d> points;
    GetSimDataInWordFrame(cameras, points);
    Eigen::Quaterniond qic(1, 0, 0, 0);
    Eigen::Vector3d tic(0, 0, 0);

    // 构建 problem
    Problem problem(Problem::ProblemType::SLAM_PROBLEM);

    // 所有 Pose
    vector<shared_ptr<VertexPose> > vertexCams_vec;
    for (size_t i = 0; i < cameras.size(); ++i) {
        shared_ptr<VertexPose> vertexCam(new VertexPose());
        Eigen::VectorXd pose(7);
        pose << cameras[i].twc, cameras[i].qwc.x(), cameras[i].qwc.y(), cameras[i].qwc.z(), cameras[i].qwc.w();
        vertexCam->SetParameters(pose);
        vertexCams_vec.push_back(vertexCam);
    }

    // add prior
//    for (size_t i = 0; i < 1; ++i)
//    {
//        shared_ptr<EdgeSE3Prior> edge(new EdgeSE3Prior(cameras[i].twc, cameras[i].qwc));
//        std::vector<std::shared_ptr<Vertex> > edge_vertex;
//        edge_vertex.push_back(vertexCams_vec[i]);  // k
//        edge->SetVertex(edge_vertex);
//        edge->SetInformation(1e4 * edge->Information());
//        problem.AddEdge(edge);
//    }

    // 所有 Point 及 edge
    std::default_random_engine generator;
    std::normal_distribution<double> noise_pdf(0, 1.);
    double noise = 0;
    vector<double> noise_invd;
    vector<shared_ptr<VertexInverseDepth> > allPoints;
    // k 控制滑动次数
    for (size_t k = 0; k < cameras.size() - window_size; ++k) {
//        int k = 0;
        if (k == 0) {
            // 第一个滑动窗口
            for (int i = 0; i < window_size; ++i) {
                problem.AddVertex(vertexCams_vec[i]);
            }
            // 添加这几帧能够看到的观测
            for (size_t i = 0; i < points.size(); ++i) {
                Eigen::Vector3d Pw = points[i].xyz;
                int start_frame_id = points[i].start_frame_id;
                int obs_cnt = points[i].obs_cnt;
                // 只加入窗口内的约束，0,...n-1 帧的所有观测
                if (start_frame_id >= window_size - 1)
                    continue;

                Eigen::Vector3d Pc = cameras[start_frame_id].Rwc.transpose() * (Pw - cameras[start_frame_id].twc);
                noise = noise_pdf(generator);
                double inverse_depth = 1 / (Pc.z() + noise);
//                double inverse_depth = 1 / (Pc.z());
                noise_invd.push_back(inverse_depth);

                // 初始化特征 vertex
                shared_ptr<VertexInverseDepth> verterxPoint(new VertexInverseDepth());
                points[i].verterx = verterxPoint;
                VecX inv_d(1);
                inv_d << inverse_depth;
                verterxPoint->SetParameters(inv_d);
                problem.AddVertex(verterxPoint);
                allPoints.push_back(verterxPoint);

                // 每个特征对应的投影误差, 第 0 帧为起始帧
                int n = std::min((int) (start_frame_id + obs_cnt), (int) (k + window_size));
                for (size_t j = start_frame_id + 1; j < n; ++j) {

                    Eigen::Vector3d pt_i = cameras[start_frame_id].featurePerId.find(points[i].featureId)->second;
                    Eigen::Vector3d pt_j = cameras[j].featurePerId.find(points[i].featureId)->second;
                    shared_ptr<EdgeReprojection> edge(new EdgeReprojection(pt_i, pt_j));
                    edge->SetTranslationImuFromCamera(qic, tic);

                    std::vector<std::shared_ptr<Vertex> > edge_vertex;
                    edge_vertex.push_back(points[i].verterx);
                    edge_vertex.push_back(vertexCams_vec[start_frame_id]);
                    edge_vertex.push_back(vertexCams_vec[j]);
                    edge->SetVertex(edge_vertex);

                    problem.AddEdge(edge);
                }
            }

        } else {
            // 添加一帧
            int idx = k + window_size - 1;

#ifdef RESET
            // 在 k 帧基础上对 frame idx pose 进行预测
            VecX param_i = vertexCams_vec[k]->Parameters();
            Qd Qi(param_i[6], param_i[3], param_i[4], param_i[5]);
            Vec3 Pi = param_i.head<3>();

            Vec3 tij = cameras[k].qwc.inverse() * (cameras[idx].twc - cameras[k].twc);
            Qd Qij = cameras[k].qwc.inverse() * cameras[idx].qwc;
            Eigen::VectorXd pose(7);
            Vec3 Pj = Qi * tij + Pi;
            Qd Qj = Qi * Qij;
            pose << Pj, Qj.x(), Qj.y(), Qj.z(), Qj.w();
            vertexCams_vec[idx]->SetParameters(pose);
#endif
            problem.AddVertex(vertexCams_vec[idx]);
            for (size_t i = 0; i < points.size(); ++i) {
                //假设所有特征点的起始帧为第0帧， 逆深度容易得到
                Eigen::Vector3d Pw = points[i].xyz;
                int start_frame_id = points[i].start_frame_id;
                int obs_cnt = points[i].obs_cnt;

                if (start_frame_id < k || start_frame_id >= idx)
                    continue;

                // 窗口倒数第二新帧上的新特征顶点
                if (start_frame_id == idx - 1) {
                    Eigen::Vector3d Pc = cameras[start_frame_id].Rwc.transpose() * (Pw - cameras[start_frame_id].twc);
                    noise = noise_pdf(generator);
                    double inverse_depth = 1 / (Pc.z() + noise);
//                    double inverse_depth = 1 / (Pc.z());
                    noise_invd.push_back(inverse_depth);
                    // 初始化特征 vertex
                    shared_ptr<VertexInverseDepth> verterxPoint(new VertexInverseDepth());
                    points[i].verterx = verterxPoint;
                    VecX inv_d(1);
                    inv_d << inverse_depth;
                    verterxPoint->SetParameters(inv_d);
                    problem.AddVertex(verterxPoint);
                    allPoints.push_back(verterxPoint);
                }

                Point3d pt = points[i];
                // 最新帧和之前所有的老帧之间的共同观测
                Eigen::Vector3d pt_i = cameras[start_frame_id].featurePerId.find(points[i].featureId)->second;
                Eigen::Vector3d pt_j = cameras[idx].featurePerId.find(points[i].featureId)->second;
                shared_ptr<EdgeReprojection> edge(new EdgeReprojection(pt_i, pt_j));
                edge->SetTranslationImuFromCamera(qic, tic);

                std::vector<std::shared_ptr<Vertex> > edge_vertex;
                edge_vertex.push_back(points[i].verterx);
                edge_vertex.push_back(vertexCams_vec[start_frame_id]);
                edge_vertex.push_back(vertexCams_vec[idx]);
                edge->SetVertex(edge_vertex);

                problem.AddEdge(edge);
            }

        }

        // add prior
//        {
//            // 上一次优化结束后的结果，作为pose prior， 固定不动。
//            VecX param_i = vertexCams_vec[k]->Parameters();
//            Qd Qi(param_i[6], param_i[3], param_i[4], param_i[5]);
//            Vec3 Pi = param_i.head<3>();
//            shared_ptr<EdgeSE3Prior> edge(new EdgeSE3Prior(Pi, Qi));
//
//            std::vector<std::shared_ptr<Vertex> > edge_vertex;
//            edge_vertex.push_back(vertexCams_vec[k]);  // k
//            edge->SetVertex(edge_vertex);
//            edge->SetInformation(1e6 * edge->Information());
//            problem.AddEdge(edge);
//        }

#ifdef RESET
        // 滑动窗口的第一帧
        VecX before = vertexCams_vec[k]->Parameters();
        Qd Qi_before(before[6], before[3], before[4], before[5]);
        Vec3 Pi_before = before.head<3>();
#endif
        problem.Solve(5);

        std::cout << "\nCompare MonoBA results after opt..." << std::endl;
        for (size_t pt_idx = 0; pt_idx < allPoints.size() ; pt_idx+=1) {
            std::cout << "after opt, point " << pt_idx << " : gt " << 1. / points[pt_idx].xyz_in_startframe.z() << " ,noise "
                      << noise_invd[pt_idx] << " ,opt " << allPoints[pt_idx]->Parameters() << std::endl;
        }
        for (int i = 0; i < vertexCams_vec.size(); ++i) {
            std::cout << "translation " << i << " :" << vertexCams_vec[i]->Parameters().head(3).transpose() << " gt: "
                      << cameras[i].twc.transpose() << std::endl;
        }
        //
#ifdef RESET
        // reset,  we should reset before hessian matrix
        VecX origin = vertexCams_vec[k]->Parameters();
        Qd Qi(origin[6], origin[3], origin[4], origin[5]);
        Vec3 Pi = origin.head<3>();
        Qd Qw2w1 = Qi * Qi_before.inverse();
        Vec3 tw2w1 = Pi - Qw2w1 * Pi_before; //cameras[k].twc
        for(int i = k; i < k + window_size; i++)
        {
            VecX param_j = vertexCams_vec[i]->Parameters();
            Qd Qj(param_j[6], param_j[3], param_j[4], param_j[5]);
            Vec3 Pj = param_j.head<3>();

            Vec3 pts_imu_j = Qw2w1.inverse() * (Pj - tw2w1);
            Qd Q = Qw2w1.inverse() * Qj;

            Eigen::VectorXd pose(7);
            pose << pts_imu_j, Q.x(), Q.y(), Q.z(), Q.w();
            std::cout << pose.transpose() << std::endl;
            vertexCams_vec[i]->SetParameters(pose);
        }
#endif
        ///          TEST MARGINALIZATION
        problem.Marginalize(vertexCams_vec[k]);
    }

    std::cout << "\nallPoint Vertex size: " << allPoints.size() << " points: " << points.size() << std::endl;
    std::cout << " allPoint size: " << allPoints.size() << std::endl;

    return 0;
}

